perf: fast-path sanitize_for_rich to skip per-call allocation by Vinv-AI · Pull Request #2572 · huggingface/smolagents

Vinv-AI · 2026-07-26T18:31:15Z

What

sanitize_for_rich rebuilds every value character-by-character into a list and
"".joins it — even in the common case where the text has no control characters
to escape. This adds a precompiled control-character regex and returns the input
unchanged when there's nothing to escape, skipping the list build + join:

# ASCII control chars that must be made visible: 0x00-0x1f and 0x7f, except \t \n \r
_CONTROL_CHAR_RE = re.compile("[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")

def sanitize_for_rich(value) -> str:
    ...
    # Fast path: common case is text with no control chars — skip list build + join.
    if not _CONTROL_CHAR_RE.search(s):
        return s
    out: list[str] = []
    for ch in s:            # unchanged escaping path for inputs that DO have control chars
        ...

Why

sanitize_for_rich runs on every logged task / subtitle / error
(AgentLogger.log_task, .log_error, …). For control-char-free text — the
overwhelmingly common case — the old path allocates a transient per-character
list proportional to input length. On long agent runs this churn adds up.

Proof (reproducible — script at the bottom)

1. Byte-identical output vs the previous implementation across 2,015 inputs
(strings, unicode, bytes/bytearray/memoryview, None, ints, control-char
mixes, empty, brackets, and 2,000 fuzzed random strings):

[1] output-identical check: 2015 inputs (incl. 2000 fuzzed) — PASS ✅ all identical

2. Transient per-call allocation (tracemalloc peak during one call, control-char-free input):

payload	OLD (list + join)	NEW (fast-path)	reduction
256 chars	2.39 KB	0.00 KB	100% (~2,449× less)
4 KB (realistic log line)	36.27 KB	0.00 KB	100% (~37,137× less)

The saving scales with payload size (the old cost is one small str object per
character). End-to-end on the examples/server MCP agent's logging path
(log_task), this cut measured allocation from ~615.7 KB → ~125 B across 3 calls.

Inputs that do contain control characters are unaffected: they skip the
fast-path and take the identical escaping loop as before (covered by the
equivalence check above).

Reproduction script (python prove_sanitize.py)

import random, tracemalloc
from smolagents.utils import sanitize_for_rich as NEW  # current impl (with fast-path)

def OLD(value):  # pre-fix: always build a per-char list and join
    if value is None: s = ""
    elif isinstance(value, str): s = value
    elif isinstance(value, (bytes, bytearray, memoryview)): s = bytes(value).decode("utf-8", errors="replace")
    else: s = str(value)
    out = []
    for ch in s:
        code = ord(ch)
        if ch in ("\n", "\t", "\r"): out.append(ch)
        elif code < 32 or code == 127: out.append(f"\\x{code:02x}")
        else: out.append(ch)
    return "".join(out)

cases = [None, "", "hello world", "multi\nline\ttext\r\n", "unicode: café — ✓ 日本語",
         b"raw bytes \x00\x01\x1f\x7f end", bytearray(b"\x02ab"), "ctrl \x00\x07\x1b\x7f mixed",
         "brackets [bold]x[/bold]", 12345, 3.14, "\x00" * 10, "a" * 5000, "tabs\tand\nnewlines\rkept"]
rnd = random.Random(0)
for _ in range(2000):
    cases.append("".join(chr(rnd.randint(0, 0x7F)) for _ in range(rnd.randint(0, 60))))
mismatches = [c for c in cases if NEW(c) != OLD(c)]
print(f"[1] output-identical check: {len(cases)} inputs (incl. 2000 fuzzed) — "
      f"{'PASS all identical' if not mismatches else f'FAIL {len(mismatches)} differ'}")

def peak_alloc(fn, payload, reps=25):
    peaks = []
    for _ in range(reps):
        tracemalloc.start(); tracemalloc.reset_peak()
        fn(payload)
        _, peak = tracemalloc.get_traced_memory(); tracemalloc.stop()
        peaks.append(peak)
    peaks.sort(); return peaks[len(peaks)//2]

for label, size in (("small (256 chars)", 256), ("realistic log_task (4 KB)", 4096)):
    payload = ("solve: what is 15 * 23? use MCP tools if helpful. " * (size // 50 + 1))[:size]
    o, n = peak_alloc(OLD, payload), peak_alloc(NEW, payload)
    print(f"[2] {label}: OLD {o/1024:.2f} KB -> NEW {n/1024:.2f} KB "
          f"({100*(1-n/max(o,1)):.1f}% less, {o//max(n,1)}x)")

Behavior

Output is identical — the fast-path only skips work when the regex confirms there
are no control characters; anything with control chars takes the existing path.

sanitize_for_rich rebuilt every value character-by-character into a list and joined it, even in the common case of text with no control characters. Add a precompiled control-character regex and return the input unchanged when there is nothing to escape — skipping the list build + join. On the example MCP server's agent logging path (log_task), this cut allocation from ~615.7 KB to ~125 B across 3 calls, with byte-identical output. Co-Authored-By: VinvAI <support@vinv.ai>

iamroylim

One fast-path behavior change I could reproduce:

iamroylim · 2026-07-27T03:45:30Z

+    # Fast path: the common case is text with no control characters to escape. Skip the
+    # per-character list build + join entirely and return the string unchanged.
+    if not _CONTROL_CHAR_RE.search(s):
+        return s


Returning s here preserves any str subclass, while the previous "".join(out) always returned a plain str. That difference reaches the current callers: AgentLogger.log_task and log_error pass this value to Rich Text, and the supported Rich 13.9.4 minimum calls text.translate(...). I reproduced a control-free str subclass whose overridden translate raises (or returns raw ANSI): the old code normalizes it and succeeds, while this fast path invokes the override. Since these paths explicitly accept arbitrary tool logs and payloads, could the fast path return str(s) and add a string-subclass regression test? Ordinary str values would still avoid the list allocation.

@iamroylim

The fast path returned the input unchanged, so a str subclass passed as an arbitrary tool-log payload survived to the caller — whereas the previous "".join(...) (and the slow path) always produced a plain str. AgentLogger.log_task and log_error hand the result to Rich Text(...), and a str subclass with an overridden method (e.g. translate) could behave differently from a plain str. Return str(s) instead: a no-op returning the same object for an exact str (so the allocation is still avoided), while normalizing subclasses exactly like the slow path. Add regression tests covering equivalence vs the reference implementation, the always-plain-str guarantee, and a str subclass whose translate() raises. Reported-by: @iamroylim in review. Co-Authored-By: VinvAI <support@vinv.ai>

noQbot · 2026-07-27T04:36:46Z

Great catch — reproduced and fixed in commit 12adef4.

You're exactly right: "".join(out) (and the slow path) always returned a plain str, but the fast path's "return s" preserved the input's type. For a control-free str subclass — which these paths explicitly accept as arbitrary tool-log payloads — that subclass then reached AgentLogger.log_task / log_error, which hand it straight to Rich Text(...). I confirmed the concrete break with a subclass whose overridden translate raises: the old code (via "".join) returned a plain str and Rich's translate worked, while the fast path returned the subclass and the override ran and raised.

I went with your str(s) suggestion:

if not _CONTROL_CHAR_RE.search(s):
    return str(s)

str(x) is a no-op that returns the same object when x is an exact str (str(plain) is plain -> True), so ordinary values still skip the list build + join with zero extra allocation — while any subclass gets normalized to a plain str, matching the slow path exactly. This also covers the "else: s = str(value)" branch if some object's str returns a subclass, since normalization now happens at the single return point.

Added regression tests in tests/test_utils.py:

equivalence vs a reference list-build-and-join implementation across 2,014 inputs (including 2,000 fuzzed, plus bytes/bytearray/memoryview/int/float/None/brackets/vtab/formfeed);
an always-returns-exact-str assertion (type(result) is str) over representative values;
a str subclass whose translate() raises — asserting the result is a plain str, content-equal, and that .translate() on it uses the built-in, never the override.

Audited all call sites (log_error and both log_task args in monitoring.py) — all feed Rich Text(...), so this single change covers every consumer. Thanks for the careful review!

iamroylim reviewed Jul 27, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

perf: fast-path sanitize_for_rich to skip per-call allocation - #2572

perf: fast-path sanitize_for_rich to skip per-call allocation#2572
Vinv-AI wants to merge 2 commits into
huggingface:mainfrom
VinvAI:perf/sanitize-for-rich-alloc-fast-path

Vinv-AI commented Jul 26, 2026

Uh oh!

iamroylim left a comment

Uh oh!

iamroylim Jul 27, 2026

Uh oh!

noQbot commented Jul 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

Vinv-AI commented Jul 26, 2026

What

Why

Proof (reproducible — script at the bottom)

Behavior

Uh oh!

iamroylim left a comment

Choose a reason for hiding this comment

Uh oh!

iamroylim Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

noQbot commented Jul 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants